package mdp.id;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

import rddl.RDDL.LCONST;
import rddl.RDDL.PVAR_NAME;

import util.Pair;
import logic.add.AADD;
import logic.add.AADDRNode;
import logic.add.ADDDNode;
import logic.add.ADDNode;
import logic.add.AADDINode;
import logic.add.DD;
import mdp.id.RDDL2DD.DEBUG;
import mdp.id.RDDL2DD.ORDER;

public class AADDHandler extends LogicRDDL2DD<AADD>{

	public AADDHandler(DEBUG dbg, ORDER ord) {
		_debug = dbg;
		_order = ord;
	}

	
	@Override
	public int approximate(int dd, double epsilon) {
		double oldprec = _context.PRECISION;
		_context.PRECISION = epsilon;
		int ret = _context.pruneNodes(dd);
		_context.PRECISION = oldprec;
		return ret;
	}

	@Override
	public int argMaxToNegInf(int v, int q) {
		int diff = _context.applyInt(v, q, DD.ARITH_MINUS);
		int actzero = thresholdAt(diff, 0.0d);
		int ret = _context.applyInt(DD_ONE, actzero, DD.ARITH_MINUS);
		return ret;
	}

	private int argMaxHelper(ADDNode node) {
		if( node instanceof ADDDNode ){
			return DD_ONE;
		}
		
		AADDINode inode = (AADDINode)node;
		ADDNode truenode = _context.getNode(inode._nHigh);
		ADDNode falsenode = _context.getNode(inode._nLow);
		int truebr = _context.getVarNode(inode._nGlobalID, 0.0d, 1.0d); 
		int falsebr = _context.getVarNode(inode._nGlobalID, 1.0d, 0.0d);
		String varname = getActualVarName(inode._nGlobalID);
		
		if(  _alactionvars.contains(varname) ){
			
			double max = inode._dMaxUpper;
			double maxt, maxf;
			
			if( truenode instanceof AADDINode ){
				maxt = ((AADDINode)truenode)._dMaxUpper;
			}else{
				maxt = ((ADDDNode)truenode)._dUpper;
			}
			
			if( falsenode instanceof AADDINode ){
				maxf = ((AADDINode)falsenode)._dMaxUpper;
			}else{
				maxf = ((ADDDNode)falsenode)._dUpper;
			}
			
			
			
			if( max == inode._dHighOffset + inode._dHighMult*maxt ){
				int truedd = argMaxHelper(truenode);
				return _context.applyInt(truedd, truebr, DD.ARITH_PROD);
			}else if( max == inode._dLowOffset + inode._dLowMult*maxf ){
				int falsedd = argMaxHelper(falsenode);
				return _context.applyInt(falsebr, falsedd, DD.ARITH_PROD);
			}else{
				System.err.println(max + " " + maxt + " " + maxf);
				System.err.println("max not equal to any....");
				try {
					System.in.read();
				} catch (IOException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
			
		}else if( _alstatevars.contains(varname) ){
			//ignore offsets
			//keep mults to 1
			//recurse
			int truedd = argMaxHelper(truenode);
			int falsedd = argMaxHelper(falsenode);
			
			truebr = _context.applyInt(truebr, truedd, DD.ARITH_PROD);
			falsebr = _context.applyInt(falsebr, falsedd, DD.ARITH_PROD);
			return _context.applyInt(falsebr, truebr, DD.ARITH_SUM);
		}else{
			System.err.println("next state var?");
			System.err.println(varname);
			try {
				System.in.read();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		
		return -1;
	}
		
		
	@Override
	protected AADD createInstance() {
		return new AADD(_alOrder);
	}

	@Override
	public int fix(int dd) {
		return _context.applyInt(dd, DD_ZERO, DD.ARITH_MAX);
	}

	@Override
	public String getActualVarName(int ddID) {
		return _tmID2VarName.get(ddID);
	}

	@Override
	public double getMax(int dd) {
		return _context.getMax(dd);
	}

	@Override
	public int getMaximalPath(int dd) {
		TreeMap<String, Boolean> actAssigns = new TreeMap<String, Boolean>();
		AADDINode nod = (AADDINode)_context.getNode(dd);
		
		ADDNode node = nod; 
		
		while( true ){
			if( node instanceof ADDDNode ){
				break;
			}else{
				
				
				AADDINode n = (AADDINode)node;
				
				int varid = n._nGlobalID;
				String var = getActualVarName(varid);

				Pair<PVAR_NAME, ArrayList<LCONST>> act = _tmActionVars.get(var);
				Boolean defaultval = (Boolean) _state.getDefaultValue(act._o1);
				
				ADDNode truthnode = _context.getNode(n._nHigh);
				ADDNode falsenode = _context.getNode(n._nLow);
				
				double max = n._dMaxUpper;
				double maxt, maxf;
				
				if( truthnode instanceof ADDDNode ) {
					maxt = ((ADDDNode)truthnode)._dUpper;
				}else {
					maxt = ((AADDINode)truthnode)._dMaxLower;
				}
				
				if( falsenode instanceof ADDDNode ) {
					maxf = ((ADDDNode)falsenode)._dUpper;
				}else {
					maxf = ((AADDINode)falsenode)._dMaxLower;
				}
				
				if( Math.abs( max - (n._dLowOffset + n._dLowMult * maxf) ) < _context.PRECISION ) {
					if( defaultval != false ) {
						actAssigns.put(var, false);
					}
					node = falsenode;
				}else if( Math.abs( max - (n._dHighOffset + n._dHighMult * maxt) ) < _context.PRECISION ){
					if( defaultval != true ) {
						actAssigns.put(var, true);
					}
					node = truthnode;
				}else {
					System.err.println("fudged up bounds!");
					System.exit(1);
				}
			}
		}
		
		Map<Integer, Boolean> intAssigns = convertToInts(actAssigns);
		
		return _context.getConstantNode(getJointActionID(intAssigns));
	}

	@Override
	public double getMin(int dd) {
		return _context.getMin(dd);
	}

	@Override
	public double getNegInf() {
		return 0;
	}

	@Override
	public Pair<AADD, TreeMap<Integer, String>> getNewContext() {
		return new Pair<AADD, TreeMap<Integer,String>>(new AADD(_alOrder), _tmID2VarName);
	}

	@Override
	public boolean getRewardActionDependent() {
		return _bRewardActionDependent;
	}

	@Override
	protected int makeNegInfDD(int dd, double val) {
		return dd;
	}

	@Override
	public int makePolicy(int v, int q) {
		int diff = _context.applyInt(v, q, DD.ARITH_MINUS);
		int actzero = thresholdAt(diff, 0.0d);
		int dd = _context.applyInt(DD_ONE, actzero, DD.ARITH_MINUS);
			
		AADDRNode refnode = _context.getRNode(dd);
		
		ArrayList<Pair<Double, Double>> affine = new ArrayList<Pair<Double,Double>>();
		affine.add(new Pair<Double, Double>(refnode._dOffset, refnode._dMult));
		
		return makePolicyHelper(refnode._nRefID,
									new HashMap<String, Boolean>(),
									affine );
	}

	@Override
	public int makePolicyHelper(int dd, Map<String, Boolean> assign,
									ArrayList<Pair<Double, Double>> affine) {
		ADDNode node = _context.getNode(dd);
		
		if( node instanceof ADDDNode ){
			double leafval = ((ADDDNode)node)._dUpper;
			
			for( int i = affine.size()-1; i >= 0; --i ) {
				Pair<Double, Double> p = affine.get(i);
				leafval = leafval*p._o2+p._o1; 
			}
			
			if( Math.abs(leafval-1) < _context.PRECISION ) {
				Map<Integer, Boolean> ints = convertToInts(assign);
				long actid = getJointActionID(ints);
				return _context.getConstantNode(actid);
			}else {
				return DD_ZERO;
			}
		}
			
		AADDINode inode = (AADDINode)node;
			
		int varid = inode._nGlobalID;
		String var = getActualVarName(varid);
		int ret;
			
		Map<String, Boolean> thisAssign = new HashMap<String, Boolean>(assign);
		
		if( _alactionvars.contains(var) ) {
			ArrayList<Pair<Double, Double>> tempaffine = new ArrayList<Pair<Double, Double>>(affine);
			tempaffine.add(new Pair<Double, Double>(inode._dHighOffset, inode._dHighMult));
			thisAssign.put(var, true);
			
			int truedd = makePolicyHelper(
							inode._nHigh, thisAssign, tempaffine);
			
			tempaffine = new ArrayList<Pair<Double, Double>>(affine);
			tempaffine.add(new Pair<Double, Double>(inode._dLowOffset, inode._dLowMult));
			thisAssign.put(var, false);
			int falsedd =  makePolicyHelper(
								inode._nLow, thisAssign, tempaffine);
				
			//places where truedd has an action
			int truepol = thresholdAt(truedd, 0.0d);
			int zeroout = _context.applyInt(DD_ONE, truepol, DD.ARITH_MINUS);
				
			ret = _context.applyInt(
					truedd, 
					_context.applyInt(falsedd, zeroout, DD.ARITH_PROD), 
				DD.ARITH_SUM);
			}else {
				int truedd = makePolicyHelper(
								inode._nHigh, thisAssign, affine);
				int falsedd =  makePolicyHelper(
								inode._nLow, thisAssign, affine);
				int truebr = _context.getVarNode(varid, 0.0d, 1.0d);
				int falsebr = _context.getVarNode(varid, 1.0d, 0.0d);
				
				ret = _context.applyInt(
						_context.applyInt(truebr, truedd, DD.ARITH_PROD), 
						_context.applyInt(falsebr, falsedd, DD.ARITH_PROD), 
					DD.ARITH_SUM);
			}
			
			return ret;
	}

	@Override
	protected void makePrimeRemap() {
		for( String s : _alstatevars ){
			_hmPrimeRemap.put(toContextID(s), toContextID(s+"'"));
		}
	}

	@Override
	public void showGraph(int dd) {
		_context.getGraph(dd, _tmID2VarName).launchViewer();
	}

	@Override
	public int swapForPrimes(int dd) {
		return _context.remapGIDsInt(dd, _hmPrimeRemap);
	}

	@Override
	public int thresholdAt(int dd, double val) {
		AADDRNode refnode = _context.getRNode(dd);
		saveThis(dd);
		Pair<Double, Double> trans = new Pair<Double, Double>(refnode._dMult, refnode._dOffset);
		ArrayList<Pair<Double, Double>> assign = new ArrayList<Pair<Double,Double>>();
		assign.add(trans);
		int ret = thresholdAtHelper(refnode._nRefID, assign, val);
//		System.out.println(threshHit);
		forgetThis(dd);
		return ret;
	}

	private int thresholdAtHelper(int refID, ArrayList<Pair<Double, Double>> assign, double thresh) {
		
		ADDNode node = _context.getNode(refID);
		if( node instanceof ADDDNode ){
			
			double leafval = ((ADDDNode)node)._dUpper;
			for( int i = assign.size()-1; i >= 0; --i ){
				Pair<Double, Double> ass = assign.get(i);
				leafval = leafval*ass._o1 + ass._o2;
			}
//			System.out.println(leafval);
			if( leafval - thresh > _context.PRECISION ){
//				System.out.println(leafval + " one");
				return DD_ONE;
			}else{
//				System.out.println(leafval + " zero");
				return DD_ZERO;
			}
			
			
		}else{
//			ArrayList<Pair<Double, Double>> thisassign = new ArrayList<Pair<Double, Double>>(assign);
			AADDINode inode = ((AADDINode)node);
			
//			if( Math.abs(inode._dMaxUpper-thresh) <= _context.PRECISION ) {
//				return DD_ZERO;
//			}
//			
//			if( Math.abs(inode._dMinLower-thresh) > _context.PRECISION ) {
//				return DD_ONE;
//			}
			
			assign.add(new Pair<Double, Double>(inode._dHighMult, inode._dHighOffset));
			int truedd = thresholdAtHelper(inode._nHigh, assign, thresh);
			int truebr = _context.getVarNode(inode._nGlobalID, 0.0d, 1.0d);
			truebr = _context.applyInt(truebr, truedd, DD.ARITH_PROD);
			assign.remove(assign.size()-1);
			
			assign.add(new Pair<Double, Double>(inode._dLowMult, inode._dLowOffset));
			int falsedd = thresholdAtHelper(inode._nLow, assign, thresh);
			int falsebr = _context.getVarNode(inode._nGlobalID, 1.0d, 0.0d);
			falsebr = _context.applyInt(falsedd, falsebr, DD.ARITH_PROD);
			assign.remove(assign.size()-1);
			
			int ret = _context.applyInt(truebr, falsebr, DD.ARITH_SUM);
			
//			flushCaches(false, ret);
			
			return ret;
		}
	}


	@Override
	public int toContextID(String varname) {
		return _tmVarName2ID.get(varname);
	}

	@Override
	protected String toContextString(String varname) {
		return String.valueOf(_tmVarName2ID.get(varname));
	}


	/* (non-Javadoc)
	 * @see mdp.id.LogicRDDL2DD#constrainThis(int, java.util.Map)
	 */
	@Override
	public int ConstrainThis(int dd, Map<Integer, Boolean> act,
			boolean stateact, boolean flush,
			final int... save) {
		
		if( _debug.compareTo(DEBUG.SOLUTION_INFO) >= 0 ) {
			System.out.println("Constraining " + getSize(dd));
			System.out.println(showMemory());
		}
		
		int ret = dd;
		for( Pair<Integer, Integer> p : _hsPairConstraints ) {
			
			int con = p._o1;
			con = isExogenous(con) ? con : simplify((Integer)p._o1, act);
			
			if( ( !stateact && isExogenous(con) ) || ( stateact ) ) {
				ret = _context.applyInt(ret, con, DD.ARITH_PROD);	
			}
			
			if( flush ) {
				saveThis(ret);
				flushCaches(false, save);
				forgetThis(ret);
			}
			
		}
		
		return ret;
	}


	/* (non-Javadoc)
	 * @see mdp.id.LogicRDDL2DD#getVarName(int)
	 */
	@Override
	public String getVarName(int pGid) {
		return _tmID2VarName.get(pGid);
	}


	/* (non-Javadoc)
	 * @see mdp.id.LogicRDDL2DD#getValueADD()
	 */
	@Override
	public int getADD(int dd, ADDHandler handl) {
		int thatdd = _context.getADD(dd, handl);
		return thatdd;
	}


	/* (non-Javadoc)
	 * @see mdp.id.LogicRDDL2DD#addPolicyConstraint(int)
	 */
	@Override
	public void addPolicyConstraint(int pPi) {
		_hsPairConstraints.add(new Pair<Integer, Integer>(pPi,-1));
	}


	/* (non-Javadoc)
	 * @see mdp.id.LogicRDDL2DD#removePolicyConstraint(int)
	 */
	@Override
	public void removePolicyConstraint(int pPi) {
		_hsPairConstraints.remove(new Pair<Integer, Integer>(pPi,-1));
	}




}